{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "from promptify.models.nlp.openai_model import OpenAI\n",
    "from promptify.prompts.nlp.prompter import Prompter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = OpenAI(api_key=\"\")\n",
    "nlp_prompter = Prompter(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = json.load(open(\"data/sql.json\",'r'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'sentence': 'Retrieve the details of all customers.', 'query': 'SELECT * FROM customers;'} \n",
      "\n",
      "{'sentence': 'Show the total number of orders placed by each customer.', 'query': 'SELECT customer_id, COUNT(*) as num_orders FROM orders GROUP BY customer_id;'} \n",
      "\n",
      "{'sentence': 'Display the total revenue generated by each product category.', 'query': 'SELECT product_name, SUM(price * quantity) as revenue FROM orders GROUP BY product_name;'} \n",
      "\n",
      "{'sentence': 'Retrieve the list of all orders placed by a specific customer.', 'query': 'SELECT * FROM orders WHERE customer_id = 123;'} \n",
      "\n",
      "{'sentence': 'Get the details of all customers whose last name starts with \"S\".', 'query': \"SELECT * FROM customers WHERE last_name LIKE 'S%';\"} \n",
      "\n"
     ]
    }
   ],
   "source": [
    "examples = []\n",
    "for sample in data[:5]:\n",
    "    print(sample,\"\\n\")\n",
    "    examples.append((sample['sentence'],sample['query']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_schema = {\n",
    "    \"customers\": {\n",
    "        \"customer_id\": int,\n",
    "        \"first_name\": str,\n",
    "        \"last_name\": str,\n",
    "        \"email\": str,\n",
    "        \"phone\": str,\n",
    "        \"address\": str\n",
    "    },\n",
    "    \"orders\": {\n",
    "        \"order_id\": int,\n",
    "        \"customer_id\": int,\n",
    "        \"order_date\": str,\n",
    "        \"product_name\": str,\n",
    "        \"price\": float,\n",
    "        \"quantity\": int\n",
    "    }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema ={\n",
    "    \"cars\":{\n",
    "        \"car_id\":int,\n",
    "        \"car_name\":str,\n",
    "        \"car_price\":float,\n",
    "        \"car_color\":str,\n",
    "        \"car_model\":str,\n",
    "        \"car_year\":int\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Writer\n",
      "You are a highly intelligent and accurate SQL query creator. You take a sentence and turn it into a SQL query. \n",
      "Sometimes you are also provided with a table and you have to create a query that returns the correct answer.\n",
      "Your output format is a dictionary with a single key 'Q' and the value is the SQL query, so [{'Q':Query}] form, no other form.{'customers': {'customer_id': <class 'int'>, 'first_name': <class 'str'>, 'last_name': <class 'str'>, 'email': <class 'str'>, 'phone': <class 'str'>, 'address': <class 'str'>}, 'orders': {'order_id': <class 'int'>, 'customer_id': <class 'int'>, 'order_date': <class 'str'>, 'product_name': <class 'str'>, 'price': <class 'float'>, 'quantity': <class 'int'>}}\n",
      "Examples:\n",
      "\n",
      "Input: Retrieve the details of all customers.\n",
      "Output: [{'Q': 'SELECT * FROM customers;' }]\n",
      "\n",
      "Input: Show the total number of orders placed by each customer.\n",
      "Output: [{'Q': 'SELECT customer_id, COUNT(*) as num_orders FROM orders GROUP BY customer_id;' }]\n",
      "\n",
      "Input: Display the total revenue generated by each product category.\n",
      "Output: [{'Q': 'SELECT product_name, SUM(price * quantity) as revenue FROM orders GROUP BY product_name;' }]\n",
      "\n",
      "Input: Retrieve the list of all orders placed by a specific customer.\n",
      "Output: [{'Q': 'SELECT * FROM orders WHERE customer_id = 123;' }]\n",
      "\n",
      "Input: Get the details of all customers whose last name starts with \"S\".\n",
      "Output: [{'Q': 'SELECT * FROM customers WHERE last_name LIKE 'S%';' }]\n",
      "\n",
      "{'cars': {'car_id': <class 'int'>, 'car_name': <class 'str'>, 'car_price': <class 'float'>, 'car_color': <class 'str'>, 'car_model': <class 'str'>, 'car_year': <class 'int'>}}\n",
      "\n",
      "For the target_schema given above interpret the Input and Output the Query\n",
      "Input: Write a SQL query to get car names and prices from the cars table where the car color is red and the car price is greater than 10000.\n",
      "Output:\n"
     ]
    }
   ],
   "source": [
    "prompt = nlp_prompter.generate_prompt('sql_writer.jinja',\n",
    "                                      examples=examples,\n",
    "                                        example_schema=example_schema,\n",
    "                                        schema = schema,\n",
    "                                      text_input=\"Write a SQL query to get car names and prices from the cars table where the car color is red and the car price is greater than 10000.\",\n",
    "                                     description=\"SQL Writer\")\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = nlp_prompter.fit('sql_writer.jinja',\n",
    "                                      examples=examples,\n",
    "                                        example_schema=example_schema,\n",
    "                                        schema = schema,\n",
    "                                      text_input=\"Write a SQL query to get car names and prices from the cars table where\\\n",
    "                                      the car color is red and the car price is greater than 10000.\",\n",
    "                                     description=\"SQL Writer\",\n",
    "                                     model_name=\"text-davinci-003\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'prompt_tokens': 558,\n",
       " 'completion_tokens': 34,\n",
       " 'total_tokens': 592,\n",
       " 'text': \" [{'Q': 'SELECT car_name, car_price FROM cars WHERE car_color = 'red' AND car_price > 10000;' }]\"}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "prompt-env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "003cbd9d5de60837b71ec248298ceb3401870867de42df8ae6f444c9dc068945"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
